﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Windows.Media.Imaging;

namespace System.Windows.Media.Media3D10
{
    /// <summary>
    /// Calcule une image de Starburst
    /// </summary>
    public class Starburst
    {
        private const double NEUTRAL_LAMBDA = 575.0E-9;

        /// <summary>
        /// OuvertureFocale de l'objectif de la camera
        /// </summary>
        public BitmapSource CameraApertureMask { get; set; }

        /// <summary>
        /// Maque de bruit
        /// </summary>
        public BitmapSource NoiseMask { get; set; }

        /// <summary>
        /// Distance focale de la l'objectif
        /// </summary>
        public double FocalDistance { get; set; }

        /// <summary>
        /// Calcule l'effet Starburst avec les paramètres de l'objet
        /// </summary>
        /// <returns>Retourne un <see cref="BitmapSource"/> qui représente l'effet calculé</returns>
        public BitmapSource Compute()
        {
            double focal = this.FocalDistance * 1E12;
            int w = CameraApertureMask.PixelWidth;
            int h = CameraApertureMask.PixelHeight;
            double dpiX = CameraApertureMask.DpiX;
            double dpiY = CameraApertureMask.DpiY;
            int stride = w * 4;

            byte[] aperturePixels = new byte[stride * h];
            CameraApertureMask.CopyPixels(aperturePixels, stride, 0);
            byte[] grayPixels = CreateGrayImage(aperturePixels);

            byte[] noisePixels = new byte[stride * h];
            NoiseMask.CopyPixels(noisePixels, stride, 0);
            byte[] grayMaskPixels = CreateGrayImage(noisePixels);

            double[] mask = new double[stride * h];
            for (int i = 0; i < stride * h; i++)
            {
                mask[i] = 1.0;
            }
            Complex[,] maskInput = CreateComplexInput(mask, w, h, stride, NEUTRAL_LAMBDA, focal);
            Color[,] maskColor = new Color[w, h];
            for (int i = 0; i < w; i++)
            {
                for (int j = 0; j < h; j++)
                {
                    byte gray = (byte)(maskInput[i, j].real * 255);
                    maskColor[i, j].R = gray;
                    maskColor[i, j].G = gray;
                    maskColor[i, j].B = gray;
                    maskColor[i, j].A = 255;
                }
            }
            CreateOutputPixels(maskColor).ToBitmapSource(w, h, stride, dpiX, dpiY, PixelFormats.Bgra32).Save("D:\\mask.png", new PngBitmapEncoder());

            double[] inputPixels = CreateInputPixels(aperturePixels, noisePixels);
            Complex[,] input = CreateComplexInput(inputPixels, w, h, stride, NEUTRAL_LAMBDA, focal);
            Color[,] inputColor = new Color[w, h];
            for (int i = 0; i < w; i++)
            {
                for (int j = 0; j < h; j++)
                {
                    byte gray = (byte)(input[i, j].real * 255);
                    inputColor[i, j].R = gray;
                    inputColor[i, j].G = gray;
                    inputColor[i, j].B = gray;
                    inputColor[i, j].A = 255;
                }
            }
            CreateOutputPixels(inputColor).ToBitmapSource(w, h, stride, dpiX, dpiY, PixelFormats.Bgra32).Save("D:\\input.png", new PngBitmapEncoder());

            Complex[,] fft = FastFourierTransform.ComputeForward(input);
            double[,] magnitudeSquared = FastFourierTransform.GetMagnitudeSquaredPlot(fft);

            double[,] magnitude = FastFourierTransform.GetMagnitudePlot(fft);
            Color[,] fftColor = new Color[w, h];
            for (int i = 0; i < w; i++)
            {
                for (int j = 0; j < h; j++)
                {
                    byte gray = (byte)(magnitude[i, j] * 255);
                    fftColor[i, j].R = gray;
                    fftColor[i, j].G = gray;
                    fftColor[i, j].B = gray;
                    fftColor[i, j].A = 255;
                }
            }
            CreateOutputPixels(fftColor).ToBitmapSource(w, h, stride, dpiX, dpiY, PixelFormats.Bgra32).Save("D:\\FFT.png", new PngBitmapEncoder());

            double[,] psf = ComputePSF(magnitudeSquared, NEUTRAL_LAMBDA, focal);
            Color[,] spectralPSF = ComputeSpectralPSF(psf);

            byte[] outputPixels = CreateOutputPixels(spectralPSF);
            BitmapSource returnValue = outputPixels.ToBitmapSource(w, h, stride, dpiX, dpiY, PixelFormats.Bgra32);

            return returnValue;
        }

        private void Save(Complex[,] complex, string path)
        {
            int w = complex.GetLength(0);
            int h = complex.GetLength(1);
            int stride = w * 4;
            double[,] reals = GetRealPart(complex);
            Color[,] colors = new Color[w, h];
            for (int i = 0; i < w; i++)
            {
                for (int j = 0; j < h; j++)
                {
                    byte gray = (byte)reals[i, j];
                    colors[i, j] = Color.FromRgb(gray, gray, gray);
                }
            }
            byte[] testOutput = CreateOutputPixels(colors);
            BitmapSource bitmap = testOutput.ToBitmapSource(w, h, stride, CameraApertureMask.DpiX, CameraApertureMask.DpiY, PixelFormats.Bgra32);
            bitmap.Save(path, new PngBitmapEncoder());
        }

        private double[,] GetRealPart(Complex[,] complex)
        {
            int w = complex.GetLength(0);
            int h = complex.GetLength(1);
            double[,] output = new double[w,h];
            for (int i = 0; i < w; i++)
            {
                for (int j = 0; j < h; j++)
                {
                    output[i,j] = complex[i,j].real;
                }
            }
            return output;
        }

        private byte[] CreateGrayImage(byte[] sourcePixels)
        {
            byte[] output = new byte[sourcePixels.Length];
            for (int i = 0; i < sourcePixels.Length; i += 4)
            {
                byte gray = (byte)((sourcePixels[i] + sourcePixels[i + 1] + sourcePixels[i + 2]) / 3);
                output[i] = gray;
                output[i + 1] = gray;
                output[i + 2] = gray;
                output[i + 3] = 255;
            }
            return output;
        }

        private double[] CreateInputPixels(byte[] aperturePixels, byte[] noisePixels)
        {
            int n = aperturePixels.Length;
            double[] output = new double[n];
            for (int i = 0; i < n; i++)
            {
                output[i] = Math.Min(aperturePixels[i], noisePixels[i]) / 255.0;
            }
            return output;
        }

        private Complex[,] CreateComplexInput(double[] inputPixels, int width, int height, int stride, double lambda, double focalDistance)
        {
            int n = inputPixels.Length;
            int step = (int)(stride / width);
            int x, y;
            double z;
            int x0 = (int)(width / 2);
            int y0 = (int)(height / 2);
            double factor = ((Math.PI)/(lambda * focalDistance));

            Complex[,] output = new Complex[width, height];
            for (int i = 0; i < width; i++)
            {
                x = i - x0;
                for (int j = 0; j < height; j++)
                {
                    y = j - y0;
                    z = factor * (x * x + y * y);
                    Complex exp = Math10.ComplexExp(1.0, z);
                    Complex value = new Complex(inputPixels[j * stride + i * 4], 0.0);
                    output[i, j] = value * exp;
                }
            }
            return output;
        }

        private double[,] ComputePSF(double[,] magnitudeSquared, double lambda, double focalDistance)
        {
            int w = magnitudeSquared.GetLength(0);
            int h = magnitudeSquared.GetLength(1);
            double factor = 1 / (lambda * focalDistance);
            double factor2 = factor * factor * 1E12;
            double[,] output = new double[w, h];
            for (int x = 0; x < w; x++)
            {
                for (int y = 0; y < h; y++)
                {
                    output[x, y] = factor2 * magnitudeSquared[x, y];
                }
            }
            return output;
        }

        private Color[,] ComputeSpectralPSF(double[,] psf)
        {
            int w = psf.GetLength(0);
            int h = psf.GetLength(1);
            int x, y, x0, y0, xi, yi;
            double coeff = NEUTRAL_LAMBDA * 1E9 / Math.Sqrt(2);
            x0 = w / 2;
            y0 = h / 2;
            CIEXYZ sl;
            CIEXYZ[,] Fs = new CIEXYZ[w, h];
            for (int lambda = 380; lambda < 770; lambda += 5)
            {
                sl = CIEXYZ.FromLambda(lambda);
                for (x = 0; x < w; x++)
                {
                    xi = (int)((x - x0) * coeff / (double)lambda);
                    for (y = 0; y < h; y++)
                    {
                        yi = (int)((y - y0) * coeff / (double)lambda);
                        if (xi * xi < x0 * x0 && yi * yi < y0 * y0)
                        {
                            Fs[x, y] += sl * psf[x0 + xi, y0 + yi];
                        }
                    }
                }
            }

            Color[,] output = new Color[w, h];
            for (x = 0; x < w; x++)
            {
                for (y = 0; y < h; y++)
                {
                    output[x, y] = Fs[x, y].ToRGB();
                }
            }
            return output;
        }

        private byte[] CreateOutputPixels(Color[,] input)
        {
            int w = input.GetLength(0);
            int h = input.GetLength(1);
            byte[] output = new byte[input.Length * 4];
            for (int i = 0; i < input.Length; i++)
            {
                Color value = (input[i % w, i / w]);
                output[4 * i] = value.B;
                output[4 * i + 1] = value.G;
                output[4 * i + 2] = value.R;
                output[4 * i + 3] = value.A;
            }
            return output;
        }
    }
}
